from cryptography import x509
from cryptography.x509 import *
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import hashes,serialization
import ipaddress
import random
import pytz
import datetime

# 保存私钥
def save_private_key(private_key, filename):
    pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption()
    )
    with open(filename, 'wb') as pem_out:
        pem_out.write(pem)

# 保存证书
def save_certificate(certificate, filename):
    pem = certificate.public_bytes(serialization.Encoding.PEM)
    with open(filename, 'wb') as pem_out:
        pem_out.write(pem)

# 生成私钥
def generate_certificate():
    # 1. 生成 CA 私钥
    ca_private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )

    # 2. 生成自签名CA证书
    # 构造证书的主体名称
    subject = issuer = Name([
        NameAttribute(NameOID.COUNTRY_NAME, u"CN"),     # 国家 (C)
        NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"GD"),       # 州/省 (ST)
        NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),             # 城市 (L)
        NameAttribute(NameOID.ORGANIZATION_NAME, u"My CA"),                 # 组织名称 (O)
        NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, u"IT Department"),  # 组织单位 (OU)
        NameAttribute(NameOID.COMMON_NAME, u"10.0.200.89"),      # 公共名称 (CN)
    ])

    # 使用时区感知的UTC时间
    not_valid_before = datetime.datetime(2024, 11, 20, 9, 16, 30, tzinfo=pytz.utc)
    not_valid_after = datetime.datetime(2025, 11, 20, 9, 16, 30, tzinfo=pytz.utc)

    # 创建证书构建器
    builder = CertificateBuilder()
    builder = builder.subject_name(subject)  # 设置证书主体
    builder = builder.issuer_name(issuer)  # 设置证书颁发者（自签名证书，颁发者与主体相同）
    builder = builder.not_valid_before(not_valid_before)  # 设置证书的起始有效期
    builder = builder.not_valid_after(not_valid_after)  # 设置证书的结束有效期
    builder = builder.serial_number(random.randint(1, 2**20))  # 设置证书的序列号（使用时间戳）
    builder = builder.public_key(ca_private_key.public_key())  # 使用私钥的公钥

    # 3. 添加 SAN 扩展，包含域名和 IP 地址
    builder = builder.add_extension(
        x509.SubjectAlternativeName([x509.IPAddress(ipaddress.IPv4Address("10.0.200.89"))]),
        critical=False
    )

    # 添加基本约束和密钥使用扩展
    basic_constraints = BasicConstraints(ca=True, path_length=None)
    builder = builder.add_extension(basic_constraints, critical=True)

    # 计算 subjectKeyIdentifier（通过公钥生成的 SHA1 散列）
    subject_key_identifier = SubjectKeyIdentifier.from_public_key(ca_private_key.public_key())
    builder = builder.add_extension(subject_key_identifier, critical=False)

    # 计算 authorityKeyIdentifier（CA证书的公钥的 subjectKeyIdentifier）
    authority_key_identifier = AuthorityKeyIdentifier.from_issuer_subject_key_identifier(
        subject_key_identifier
    )
    builder = builder.add_extension(authority_key_identifier, critical=False)

    # 签署证书
    certificate = builder.sign(ca_private_key, hashes.SHA256(), default_backend())

    # 3. 生成服务器私钥
    server_private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=default_backend()
    )

    # 3. 生成服务器证书签名请求（CSR）
    # 设置服务器证书的主体信息
    subject = Name([
        NameAttribute(NameOID.COUNTRY_NAME, u"CN"),  # 国家
        NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, u"GD"),  # 州/省
        NameAttribute(NameOID.LOCALITY_NAME, u"San Francisco"),  # 城市
        NameAttribute(NameOID.ORGANIZATION_NAME, u"Example Server"),  # 组织名称
        NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, u"IT Department"),  # 组织单位 (OU)
        NameAttribute(NameOID.COMMON_NAME, u"10.0.200.89"),  # 公共名称（CN），通常是域名
    ])

    # 创建服务器证书签名请求
    csr_builder = CertificateSigningRequestBuilder()
    csr_builder = csr_builder.subject_name(subject)
    csr_builder = csr_builder.add_extension(
        x509.SubjectAlternativeName([x509.IPAddress(ipaddress.IPv4Address("10.0.200.89"))]),
        critical=False
    )

    # 使用服务器私钥对 CSR 进行签名
    csr = csr_builder.sign(server_private_key, hashes.SHA256(), default_backend())

    # 3. 使用CA签署服务器证书
    # 使用CA私钥和CA证书签署服务器证书
    builder = CertificateBuilder()
    builder = builder.subject_name(csr.subject)
    builder = builder.issuer_name(certificate.subject)
    builder = builder.not_valid_before(not_valid_before)  # 设置证书的起始有效期
    builder = builder.not_valid_after(not_valid_after)  # 设置证书的结束有效期
    builder = builder.serial_number(random.randint(1, 2**20))  # 使用当前时间戳作为序列号
    builder = builder.public_key(csr.public_key())  # 使用CSR中的公钥

    # 添加授权密钥标识符扩展
    builder = builder.add_extension(
        AuthorityKeyIdentifier.from_issuer_subject_key_identifier(
            subject_key_identifier
        ),
        critical=False
    )

    # 使用CA的私钥签署证书
    server_certificate = builder.sign(ca_private_key, hashes.SHA256(), default_backend())

    # 5. 生成并保存CA证书和私钥
    save_private_key(ca_private_key, "ca.key")
    save_certificate(certificate, "ca.crt")

    # 6. 生成并保存服务器证书和私钥
    save_private_key(server_private_key, "server.key")

    # 使用CA签署服务器证书
    save_certificate(server_certificate, "server.crt")

    print("服务器证书和私钥已生成并保存。")

generate_certificate()
